from langgraph.graph import StateGraph, START, END
from mcp.server import FastMCP

from insight_agent.mcp_server_graph.src.mcp_server_graph.utils.text2cypher_service import *
from insight_agent.mcp_server_graph.src.mcp_server_graph.utils.http_post import *
import logging

load_dotenv()



async def text2cypher_parallel(user_question):
    # 初始化 StateGraph，状态 schema 指明了所用状态类型
    graph = StateGraph(state_schema=CypherState)

    """ Langgraph 工作流构建 """

    # 添加节点与边
    graph.add_node("PrepareCypher", prepare_cypher)
    # graph.add_node("GenerateCypher", generate_cypher)
    graph.add_node("SampleNeo4j", sample_neo4j)
    graph.add_node("QueryNeo4j", query_neo4j)
    graph.add_node("DirectSearch", direct_search)
    graph.add_node("FixCypher", fix_cypher)
    graph.add_node("InterpretResult", interpret_result)
    
    graph.add_node("ParallelProcessing", parallel_processing)
    graph.add_node("FixCypherParallel", fix_cypher_parallel)

    # 添加入口与常规边
    graph.add_edge(START, "PrepareCypher")
    # graph.add_edge("SampleNeo4j", "GenerateCypher")
    # graph.add_edge("GenerateCypher", "QueryNeo4j")
    graph.add_edge("SampleNeo4j", "ParallelProcessing")

    # 添加条件边
    graph.add_conditional_edges("QueryNeo4j", after_query_decide_next)
    graph.add_conditional_edges("PrepareCypher", after_prepare_decide_next)
    graph.add_conditional_edges("ParallelProcessing", after_query_decide_next_parallel)
    graph.add_conditional_edges("FixCypherParallel", after_query_decide_next_parallel)

    # 循环流程与结束路径
    # graph.add_edge("FixCypher", "QueryNeo4j")
    graph.add_edge("InterpretResult", END)
    graph.add_edge("DirectSearch", "QueryNeo4j")

    """流程执行"""

    # 编译后执行
    compiled = graph.compile()
    print(f"用户问题：{user_question}")
    initial_state = {"user_question": user_question,
                     "scenario": None,  # 业务code
                     "selectSchema": None,  # 挑选的schema
                     "sampleCypher": None,  # 采样语句
                     "preCypher": [],  # 历史语句
                     "preCypherResult": [],  # 历史结果
                     "sampleReason": None,  # 业务和采样思路
                     "emptyReason": [],  # 空数据分析思路
                     "cypher": None,
                     "reason": None,
                     "query_result": None,
                     "error": None,
                     "retry_count": 0,
                     "max_retries": 10,
                     "interpretation": None,
                     "parallel_count": 3
                     }
    final = await compiled.ainvoke(initial_state,
                                   config={"recursion_limit": 60})
    print("流程完全结束")
    return final

async def text2cypher(user_question):
    # 初始化 StateGraph，状态 schema 指明了所用状态类型
    graph = StateGraph(state_schema=CypherState)

    """ Langgraph 工作流构建 """

    # 添加节点与边
    graph.add_node("PrepareCypher", prepare_cypher)
    graph.add_node("GenerateCypher", generate_cypher)
    graph.add_node("SampleNeo4j", sample_neo4j)
    graph.add_node("QueryNeo4j", query_neo4j)
    graph.add_node("DirectSearch", direct_search)
    graph.add_node("FixCypher", fix_cypher)
    graph.add_node("InterpretResult", interpret_result)
    
    # graph.add_node("ParallelProcessing", parallel_processing)

    # 添加入口与常规边
    graph.add_edge(START, "PrepareCypher")
    graph.add_edge("SampleNeo4j", "GenerateCypher")
    graph.add_edge("GenerateCypher", "QueryNeo4j")
    # graph.add_edge("SampleNeo4j", "ParallelProcessing")

    # 添加条件边
    graph.add_conditional_edges("QueryNeo4j", after_query_decide_next)
    graph.add_conditional_edges("PrepareCypher", after_prepare_decide_next)
    # graph.add_conditional_edges("ParallelProcessing", after_query_decide_next)

    # 循环流程与结束路径
    graph.add_edge("FixCypher", "QueryNeo4j")
    graph.add_edge("InterpretResult", END)
    graph.add_edge("DirectSearch", "QueryNeo4j")

    """流程执行"""

    # 编译后执行
    compiled = graph.compile()
    print(f"用户问题：{user_question}")
    initial_state = {"user_question": user_question,
                     "scenario": None,  # 业务code
                     "selectSchema": None,  # 挑选的schema
                     "sampleCypher": None,  # 采样语句
                     "preCypher": [],  # 历史语句
                     "preCypherResult": [],  # 历史结果
                     "sampleReason": None,  # 业务和采样思路
                     "emptyReason": [],  # 空数据分析思路
                     "cypher": None,
                     "reason": None,
                     "query_result": None,
                     "error": None,
                     "retry_count": 0,
                     "max_retries": 10,
                     "interpretation": None,
                     "parallel_count": 1
                     }
    final = await compiled.ainvoke(initial_state,
                                   config={"recursion_limit": 60})
    print("流程完全结束")
    return final


# 初始化 MCP 服务器
mcp = FastMCP("GraphServer",
              port=8181,
              )


@mcp.tool()
async def agent_graph(user_request=""):
    """
    任何用户问题都建议使用此工具，此工具和其他工具没有排斥关系，有很强的互补作用，建议并行调用。
    使用以自然人为中心构建的图数据，根据用户问题进行图查询，当提到人车家、异网拉新、低资费用户存量运营、诈骗识别、用户细分、使用图数据库、使用图工具这些词时，务必调用此工具，返回某场景下的用户子图结构和解读
    业务场景解释
    0. 其他问题：使用text2cypher技术生成图查询语句并查询
    1. 人车家：基于用户对4S店的访问和懂车帝APP的使用，判断用户可能存在购车需求，并根据用户家庭人数推荐车型
    2. 异网拉新：基于用户白天常驻基站和通信关系，聚合用户为地理圈群体，我们判断同一个地理圈的用户有类似的需求，我们基于每个地理圈的资费和APP使用偏好来制定圈内异网用户的拉新策略。异网用户指的是运营商非移动的用户，拉新指的是我们将异网用户推荐成为移动的新客户。
    3. 低资费用户存量运营：基于地理圈群体和群体的资费和APP情况，我们判断同一个地理圈的用户有类似的需求，对群体内资费等偏低的用户进行产品营销推荐
    4. 诈骗识别：当一个新注册号码有大量呼出并且是漫游状态，很可能具有电诈风险
    5. 用户细分：查询集团用户群体中个人订购较多的产品，并发掘产品的潜在合作企业，制定推荐策略
    6. 家庭圈查询：查询家庭成员的互联情况
    7. 诈骗易感人群：查询容易被诈骗的对象
    8. 根据企业的竞品使用信息，提供自有产品替代方案
    9. 根据新闻/国家要求发现产品推广商机
    10. 将产品推荐给同行业还未订购的其他企业
    11. 根据企业的供应链信息，将产品推荐给企业的上游
    12. 向企业推荐服务升级解决方案
    常见问题有：
    "帮我找出北京地区经常去4S店，可能有购车意向的用户" ,
    "基于海淀区校园用户的特点，进行异网用户拉新",
    "日间常驻海淀区科技公司附近的用户群体中，哪些人资费明显低于其他人，请给出营销建议"
    "帮我筛查最近一周有大量短时长国际通话的新入网用户，识别诈骗风险"
    "查询集团用户群体中个人订购较多的产品，并发掘产品的潜在合作企业，制定推荐策略"
    "查询用户U02369的二度关系网络，识别高价值用户，并给出营销策略"
    :param user_request: 用户提出的问题
    :return: 图查询结果的解读
    """
    final_state = await text2cypher_parallel(user_request)
    # final_state = await text2cypher(user_request)
    vertex_set, edge_set, other = final_state["query_result"]
    data = {"vertexSet": vertex_set, "edgeSet": edge_set}
    summary = final_state["interpretation"]

    # data 字段处理
    data["resultSet"] = {
        "header": [],
        "table": []
    }
    data["errors"] = []
    data["groupSet"] = {}
    data["profile"] = {}
    data["vertexCount"] = len(data["vertexSet"])
    data["edgeCount"] = len(data["edgeSet"])
    data["resultCount"] = 1

    # summary处理
    summary = summary.replace('markdown', '')
    summary = summary.replace('```', '')

    answer = {"data": data, "summary": summary + f"\n\n生成的cypher语句：\n```\n{final_state["cypher"]}\n```"}
    logging.info("agent_graph正在发送数据")

    update_insight_result(agent_type="agent_graph", updated_result=answer)
    logging.info("agent_graph发送数据结束")
    return summary


if __name__ == "__main__":
    # 以标准 I/O 方式运行 MCP 服务器
    mcp.run(transport='streamable-http')
